# update_db.py
# 智能数据库结构更新脚本
import sys
import os

# 添加backend目录到Python路径
backend_dir = os.path.dirname(os.path.abspath(__file__))
if backend_dir not in sys.path:
    sys.path.append(backend_dir)

from sqlalchemy import text, inspect
from sqlalchemy.exc import SQLAlchemyError

# 使用绝对导入
from backend.database.session import SessionLocal, engine, Base
from backend.models.models import User, FollowUpRecord, Diagnosis, Advice

from backend.models.models import Base, engine


# 要检查的模型列表
MODELS = [User, FollowUpRecord, Diagnosis, Advice]

def get_table_columns(table_name):
    """获取数据库表的列信息"""
    inspector = inspect(engine)
    return {col['name']: col for col in inspector.get_columns(table_name)}


def update_database_structure():
    """更新数据库结构，确保与模型定义一致"""
    db = SessionLocal()
    try:
        # 创建所有表（如果不存在）
        Base.metadata.create_all(bind=engine)
        print("确保所有表已创建")

        # 检查每个表的列结构
        for model in MODELS:
            table_name = model.__tablename__
            print(f"\n检查表: {table_name}")

            # 获取模型定义的列
            model_columns = {col.name: col for col in model.__table__.columns}
            # 获取数据库中的列
            db_columns = get_table_columns(table_name)

            # 检查缺失的列
            for col_name, col in model_columns.items():
                if col_name not in db_columns:
                    # 添加缺失的列
                    try:
                        # 构建添加列的SQL语句
                        # 注意：这里简化处理，实际可能需要更复杂的类型映射
                        col_type = str(col.type)
                        nullable = "NOT NULL" if not col.nullable else ""
                        default = f"DEFAULT '{col.default.arg}'" if col.default else ""
                        sql = f"ALTER TABLE {table_name} ADD COLUMN {col_name} {col_type} {nullable} {default}"
                        sql = sql.strip()  # 移除多余空格
                        print(f"执行SQL: {sql}")
                        db.execute(text(sql))
                        db.commit()
                        print(f"添加列成功: {table_name}.{col_name}")
                    except SQLAlchemyError as e:
                        print(f"添加列失败: {table_name}.{col_name}, 错误: {e}")
                        db.rollback()
                else:
                    # 检查列类型是否一致（简化版）
                    db_col_type = str(db_columns[col_name]['type'])
                    model_col_type = str(col.type)
                    # 粗略比较类型
                    if db_col_type.lower() != model_col_type.lower():
                        print(f"警告: 列类型不匹配 {table_name}.{col_name}: 数据库={db_col_type}, 模型={model_col_type}")
                        print("注意: SQLite不支持修改列类型，需要手动处理")

            # 检查多余的列
            for col_name in db_columns.keys():
                if col_name not in model_columns and col_name != "id":  # 假设id是主键，不删除
                    print(f"警告: 发现多余列 {table_name}.{col_name}")
                    # 对于SQLite，删除列需要复杂的操作，这里仅提示
                    print("注意: SQLite不支持直接删除列，如需删除请手动处理")

        # 特殊处理：移除follow_up_records表中的android_data列（如果存在）
        if 'follow_up_records' in inspect(engine).get_table_names():
            follow_up_columns = get_table_columns('follow_up_records')
            if 'android_data' in follow_up_columns:
                print("\n发现follow_up_records表中有android_data列，准备移除...")
                try:
                    # SQLite删除列的复杂操作
                    db.execute(text("""
                        CREATE TABLE follow_up_records_new (
                            id INTEGER PRIMARY KEY, 
                            username VARCHAR(50) NOT NULL, 
                            doctor_name VARCHAR(50) NOT NULL, 
                            date DATE NOT NULL, 
                            left_eye_image VARCHAR(255), 
                            right_eye_image VARCHAR(255),
                            FOREIGN KEY(username) REFERENCES users(username)
                        )
                    """))
                    db.execute(text("""
                        INSERT INTO follow_up_records_new (id, username, doctor_name, date, left_eye_image, right_eye_image)
                        SELECT id, username, doctor_name, date, left_eye_image, right_eye_image FROM follow_up_records
                    """))
                    db.execute(text("DROP TABLE follow_up_records"))
                    db.execute(text("ALTER TABLE follow_up_records_new RENAME TO follow_up_records"))
                    db.commit()
                    print("成功移除follow_up_records表中的android_data列")
                except SQLAlchemyError as e:
                    print(f"移除android_data列失败: {e}")
                    db.rollback()

        print("\n数据库结构更新完成！")
    except Exception as e:
        print(f"更新数据库结构时出错: {e}")
    finally:
        db.close()

def set_user_role(username: str, role: str):
    """设置用户角色"""
    db = SessionLocal()
    try:
        user = db.query(User).filter(User.username == username).first()
        if not user:
            print(f"用户 {username} 不存在")
            return False
        
        user.role = role
        db.commit()
        print(f"用户 {username} 的角色已设置为 {role}")
        return True
    finally:
        db.close()


def update_database():
    # 删除所有表
    Base.metadata.drop_all(bind=engine)
    # 创建新表
    Base.metadata.create_all(bind=engine)
    print("数据库已更新")

if __name__ == "__main__":
    update_database()
    print("\n设置管理员用户角色...")
    admin_username = 'admin'
    set_user_role(admin_username, 'admin')
    
    print("\n操作完成！")